from rest_framework import viewsets, status
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework.permissions import IsAuthenticated
from rest_framework import viewsets, filters
from django.core.files.storage import FileSystemStorage
from django.db import transaction
from django.conf import settings
from .models import KnowledgeGraph, UnstructuredFile, GraphDataSource, GraphDataSourceTable, KnowledgeGraphStatus
from rest_framework.pagination import PageNumberPagination
from .serializers import (
    KnowledgeGraphListSerializer,
    KnowledgeGraphCreateSerializer,
    UnstructuredFileSerializer,
    TestConnectionSerializer,
    KnowledgeGraphCreateSerializer,
    GraphDataSourceSerializer,
    GraphDataSourceTableSerializer,
    GraphDataSourceTableCheckSerializer
)
from .tasks import process_unstructured_file, build_structured_graph
from .db_connect import MySQLConnectionStrategy, PostgreSQLConnectionStrategy, DatabaseService
from graphNeo4j.lib.Neo4jServices import Neo4jServices
from dify.api import upload_file as dify_upload_file, run_workflow as dify_run_workflow, parse_workflow_response
from .services.patent_source_service import PatentSourceService
import logging
import requests
import json

logger = logging.getLogger('celery_tasks')


class CustomPagination(PageNumberPagination):
    """自定义分页器"""
    page_size = 10
    page_size_query_param = 'page_size'
    max_page_size = 100


class KnowledgeGraphViewSet(viewsets.ModelViewSet):
    """
    知识图谱视图集
    提供图谱的CRUD操作
    """
    pagination_class = CustomPagination
    permission_classes = [IsAuthenticated]
    queryset = KnowledgeGraph.objects.all()
    
    def get_serializer_class(self):
        if self.action == 'create':
            return KnowledgeGraphCreateSerializer
        return KnowledgeGraphListSerializer

    def get_queryset(self):
        return self.queryset.filter(created_by=self.request.user).exclude(status=KnowledgeGraphStatus.DELETED.value)

    def perform_create(self, serializer):
        graph = serializer.save(created_by=self.request.user)
        if graph.graph_type == 'structured':
            # 注释掉异步调用
            # build_structured_graph.delay(graph.id)
            pass
        return graph

    # 重写list方法
    def list(self, request, *args, **kwargs):
        """获取图谱列表"""
        queryset = self.filter_queryset(self.get_queryset())
        name = request.query_params.get('name', None)
        if name:
            queryset = queryset.filter(name__icontains=name)
        page = self.paginate_queryset(queryset)
        if not page:
            return Response({
                'code': 200,
                'message': '获取成功',
                'data': {
                    'count': 0,
                    'next': None,
                    'previous': None,
                    'results': [],
                    'empty_tip': '当前列表为空'
                }
            })

        serializer = self.get_serializer(page, many=True)
        response = self.get_paginated_response(serializer.data)

        return Response({
            'code': 200,
            'message': '获取成功',
            'data': response.data
        })

    @action(detail=False, methods=['post'])
    def generate(self, request):
        """新建图谱"""
        try:
            serializer = KnowledgeGraphCreateSerializer(data=request.data)
            serializer.is_valid(raise_exception=True)  # 这将自动返回400响应给客户端

            with transaction.atomic():  # 使用事务保证一致性
                graph = serializer.save(created_by=request.user)  # 利用序列化器的save方法，它会调用create

            return Response({
                'code': 200,
                'message': '新建图谱成功',
                'data': KnowledgeGraphListSerializer(graph).data
            })

        except Exception as e:
            return Response({
                'code': 500,
                'message': f'新建图谱失败: {str(e)}'
            }, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    @action(detail=True, methods=['post'])
    def upload_file(self, request, pk=None):
        """
        上传非结构化文件
        """
        graph = self.get_object()
        if graph.graph_type != 'unstructured':
            return Response(
                {'error': '只能为非结构化图谱上传文件'}, 
                status=status.HTTP_400_BAD_REQUEST
            )

        file_obj = request.FILES.get('file')
        if not file_obj:
            return Response(
                {'error': '未找到上传文件'}, 
                status=status.HTTP_400_BAD_REQUEST
            )

        # 保存文件
        fs = FileSystemStorage()
        filename = fs.save(file_obj.name, file_obj)
        file_path = fs.path(filename)

        # 创建文件记录
        unstructured_file = UnstructuredFile.objects.create(
            graph=graph,
            file_name=file_obj.name,
            file_path=file_path,
            file_type=file_obj.content_type
        )

        return Response({
            'code': 200,
            'message': "上传成功",
            'data': UnstructuredFileSerializer(unstructured_file).data
        })

    @action(detail=False, methods=['post'])
    def test_connection(self, request):
        """
        测试数据库连接
        """
        serializer = TestConnectionSerializer(data=request.data)
        if not serializer.is_valid():
            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)

        validated_data = serializer.validated_data
        data_source = validated_data['data_source']

        try:
            # 根据数据库类型选择相应的策略
            db_service = DatabaseService(data_source['db_type'])
            db_service.connection(data_source)
            result = db_service.test_database_connection()

            return Response({
                'code': 200,
                'message': result['message']
            })
        except Exception as e:
            return Response(
                {'error': f"连接失败！{e}"},
                status=status.HTTP_400_BAD_REQUEST
            )

    @action(detail=True, methods=['put'])
    def data_source(self, request, pk=None):
        """
        更新或创建图谱的数据源
        """
        try:
            graph = self.get_object()  # 使用DRF提供的get_object方法来获取图谱实例
            # 仅结构化图谱可以添加数据源
            if graph.graph_type != 'structured':
                return Response(
                    {'error': "非结构化数据图谱不允许添加数据源"},
                    status=status.HTTP_400_BAD_REQUEST
                )

            # 数据参数校验
            serializer = GraphDataSourceSerializer(data=request.data)
            if not serializer.is_valid():
                return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)

            # data_source = GraphDataSource.objects.get(graph=graph)
            # 尝试获取或创建 GraphDataSource 实例
            data_source, created = GraphDataSource.objects.get_or_create(
                graph=graph,
                defaults=serializer.validated_data
            )
            # 如果是更新操作，则应用新的数据
            if not created:
                for key, value in serializer.validated_data.items():
                    setattr(data_source, key, value)
                data_source.save()

            return Response({
                'code': 200,
                'message': '添加成功',
            }, status=status.HTTP_201_CREATED)
        except KnowledgeGraph.DoesNotExist:
            return Response(
                {'error': "未找到指定的图谱"},
                status=status.HTTP_404_NOT_FOUND
            )


    @action(detail=True, methods=['get'])
    def data_source_tables(self, request, pk=None):
        """
        获取数据源的表列表
        """
        try:
            graph = self.get_object()  # 使用DRF提供的get_object方法来获取图谱实例
            # 仅结构化图谱可以添加数据源
            if graph.graph_type != 'structured':
                return Response(
                    {'error': "非结构化数据图谱无法获取数据表"},
                    status=status.HTTP_400_BAD_REQUEST
                )
            data_source = getattr(graph, 'data_source', None)
            data_source = GraphDataSourceSerializer(data_source).data

            # 数据库连接
            if data_source.get('db_type') == 'mysql':
                table_query = """
                        SELECT TABLE_NAME, TABLE_SCHEMA, TABLE_TYPE, TABLE_COMMENT
                        FROM information_schema.TABLES
                        WHERE TABLE_SCHEMA = %s
                        ORDER BY TABLE_NAME
                        LIMIT %s OFFSET %s;
                    """
                total_count_query = "SELECT COUNT(*) FROM information_schema.tables WHERE TABLE_SCHEMA = %s;"
            elif data_source.get('db_type') == 'postgresql':

                table_query = """
                        SELECT c.relname AS table_name, n.nspname AS table_schema, 
                               CASE WHEN relkind = 'r' THEN 'BASE TABLE' ELSE relkind END AS table_type,
                               obj_description(c.oid) AS table_comment
                        FROM pg_class c
                        JOIN pg_namespace n ON n.oid = c.relnamespace
                        WHERE n.nspname NOT IN ('pg_catalog', 'information_schema') AND c.relkind IN ('r', 'v', 'm', 'f', 'p')
                        ORDER BY c.relname
                        LIMIT %s OFFSET %s;
                    """
                total_count_query = "SELECT COUNT(*) FROM information_schema.tables WHERE table_catalog = %s;"
            else:
                return Response(
                    {'error': '不支持的数据库类型'},
                    status=status.HTTP_400_BAD_REQUEST
                )



            db_service = DatabaseService(data_source['db_type'])
            connection = db_service.connection(data_source)

            # 获取分页参数
            page = request.query_params.get('page', 1)
            page_size = request.query_params.get('page_size', 10)

            try:
                page = int(page)
                page_size = int(page_size)
            except ValueError:
                return Response(
                    {'error': '无效的分页参数'},
                    status=status.HTTP_400_BAD_REQUEST
                )

            offset = (page - 1) * page_size

            with connection.cursor() as cursor:
                # 计算总记录数以确定总页数
                cursor.execute(total_count_query, [data_source.get('db_name')])
                total_count = cursor.fetchone()[0]
                # 执行分页查询
                cursor.execute(table_query, [data_source.get('db_name'), page_size, offset])
                tables = cursor.fetchall()

            # 将结果转换为字典列表
            table_data = [{'table_name': row[0], 'table_commit': row[3]} for row in tables]

            return Response({
                'code': 200,
                'message': '数据源表获取成功',
                'data': {
                    'count': total_count,
                    'results': table_data
                }
            })
        except KnowledgeGraph.DoesNotExist:
            return Response(
                {'error': "未找到指定的图谱"},
                status=status.HTTP_404_NOT_FOUND
            )

    @action(detail=True, methods=['post'])
    def build(self, request, pk=None):
        """
        添加要构建的数据表
        """
        try:
            graph = self.get_object()  # 使用DRF提供的get_object方法来获取图谱实例
            # 仅结构化图谱可以添加数据源
            if graph.graph_type == 'structured':
                # 数据参数校验
                serializer = GraphDataSourceTableCheckSerializer(data=request.data)
                if not serializer.is_valid():
                    return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)

                table_names = serializer.validated_data.pop('table_name', [])
                commits = serializer.validated_data.pop('commit', [])
                created_tables = []

                data_source = getattr(graph, 'data_source', None)
                data_source = GraphDataSourceSerializer(data_source).data

                # 根据数据库类型选择相应的策略
                db_service = DatabaseService(data_source['db_type'])
                db_service.connection(data_source)

                total_rows = 0
                with transaction.atomic():  # 确保所有操作要么全部成功，要么全部失败
                    for index, table_name in enumerate(table_names):
                        # 获取表中的数据条数
                        row_number = db_service.get_table_rows(table_name)
                        total_rows += row_number
                        existing_table = GraphDataSourceTable.objects.filter(graph=graph, table_name=table_name).first()
                        if existing_table:
                            return Response(
                                {'code': 400, 'message': f"Data table '{table_name}' already exists for this graph.",
                                 'data': []},
                                status=status.HTTP_400_BAD_REQUEST
                            )
                        data = {
                            'table_name': table_name,
                            'graph': graph.id,
                        }
                        if index in commits:
                            data['commit'] = commits[index]

                        table_serializer = GraphDataSourceTableSerializer(data=data)
                        if not table_serializer.is_valid():
                            return Response({
                                'code': 400,
                                'message': table_serializer.errors,
                            }, status=status.HTTP_400_BAD_REQUEST)

                        table_serializer.save()
                        created_tables.append(table_serializer.data)

                # 保存数据处理条数
                graph.build_total = total_rows
                graph.save()

                # 异步处理构建逻辑
                build_structured_graph.delay(graph.id)
            elif graph.graph_type == 'unstructured':
                total_rows = 0
                # 非结构化图谱处理逻辑
                process_unstructured_file.delay(graph.id)
            else:
                return Response(
                    {'error': "不支持的图谱类型"},
                    status=status.HTTP_400_BAD_REQUEST
                )

            return Response({
                'code': 200,
                'message': '构建成功',
                'already_build': 0,
                'build_total': total_rows
            }, status=status.HTTP_201_CREATED)

        except KnowledgeGraph.DoesNotExist:
            return Response(
                {'error': "未找到指定的图谱"},
                status=status.HTTP_404_NOT_FOUND
            )

    @action(detail=True, methods=['get'])
    def progress(self, request, pk=None):
        try:
            graph = self.get_object()

            return Response({
                'code': 200,
                'message': '获取成功',
                'data': {
                    'status': graph.status,
                    'already_build': graph.already_build,
                    'build_total': graph.build_total,
                }
            })
        except KnowledgeGraph.DoesNotExist:
            return Response(
                {'error': "未找到指定的图谱"},
                status=status.HTTP_404_NOT_FOUND
            )

    @action(detail=True, methods=['get'])
    def info(self, request, pk=None):
        try:
            graph = self.get_object()

            # 获取neo4j 数据库配置
            neo4j_config = settings.NEO4J_CONFIG
            neo4j_services = Neo4jServices(neo4j_config['uri'], neo4j_config['user'], neo4j_config['password'])
            # node_result = neo4j_services.add_knowledge_to_graph(graph_id, row)
            # 获取节点数量
            node_number = neo4j_services.get_knowledge_graph_node_number(graph.id)
            # 获取节点下所有关系数量
            relation_number = neo4j_services.get_knowledge_graph_relation_number(graph.name)
            # 获取节点列表
            nodes, relationships = neo4j_services.get_knowledge_graph_relationships(graph.name, graph.id, -1, 500)

            # 数据去重
            # 使用字典保持唯一性
            unique_data = {}
            del_node_ids = []
            for item in nodes:
                # 如果名称不在字典中，则添加
                if item['name'] not in unique_data:
                    unique_data[item['name']] = item
                else:
                    del_node_ids.append(item['nodeId'])

            # 将字典的值转换为列表
            unique_data_list = list(unique_data.values())

            # for item in relationships:
            #     if item['sourceID'] in del_node_ids or item['targetID'] in del_node_ids:
            #         relationships.remove(item)
            # 获取关系数量
            return Response({
                'code': 200,
                'message': '获取成功',
                'data': {
                    'id': graph.id,
                    'name': graph.name,
                    'graph_type': graph.graph_type,
                    'node_number':node_number,
                    'relation': relation_number,
                    'links': {
                        'nodes': unique_data_list,
                        'relationships': relationships
                    }
                }
            })
        except KnowledgeGraph.DoesNotExist:
            return Response(
                {'error': "未找到指定的图谱"},
                status=status.HTTP_404_NOT_FOUND
            )

    @action(detail=False, methods=['get'])
    def test_add_neo4j_node(self, request):

        # 查询结构化图谱信息
        graph = KnowledgeGraph.objects.get(id=2)

        data_source = getattr(graph, 'data_source', None)
        data_source = GraphDataSourceSerializer(data_source).data

        data_source_tables = getattr(graph, 'data_source_table', None)
        data_source_tables = GraphDataSourceTableSerializer(data_source_tables, many=True).data

        # 数据库连接
        db_service = DatabaseService(data_source['db_type'])
        connection = db_service.connection(data_source)

        def loop_write_data(table_name):
            page = 1
            while(1):
                if page == 3:
                    break
                data = db_service.get_table_data(table_name, page)
                for item in data:
                    row = PatentSourceService().get_patent_standard_field_data(table_name, item)
                    NEO4J_CONFIG = {
                        # "uri": "neo4j://119.254.155.85:37687",
                        "uri": "neo4j://127.0.0.1:37686",
                        "user": "neo4j",
                        "password": "nbc98765"
                        # "password": "vPn**12@@56"
                    }
                    neo4j_services = Neo4jServices(NEO4J_CONFIG['uri'], NEO4J_CONFIG['user'], NEO4J_CONFIG['password'])

                    neo4j_services.add_knowledge_to_graph(row)
                if not data:
                    break
                page += 1

        tables = []
        logger.info(f"处理数据")
        if data_source_tables:
            for row in data_source_tables:
                table_name = row['table_name']
                tables.append({
                    "table_name": table_name
                })
                loop_write_data(table_name)

        # data = PatentSourceService().get_patent_standard_field_data('cyl_patent_new_2_json', {
        #     "patent_name": "一种功能复合材料原位电控增材制造装置及工艺",
        #     "patent_number": "CN202210550306.9",
        #     'province': '陕西省',  # 省份
        #     'year': '2023',  # 年份
        #     'applicant_institution': '西安交通大学',  # 专利申请人
        #     'application_number': 'CN202210550306.9',  # 专利申请号
        #     'application_date': '2022-05-20T00:00:00Z',  # 申请（专利）日期
        #     'patent_category': '发明授权',  # 专利分类
        #     'abstract': '陕西省西安市咸宁西路28号',  # 概要
        #     'keywords': '["高压静电发生器","电压值","导线连接","增强材料","基体材料","功能复合材料","增材制造"]',  # 关键词
        # })
        #


        return Response({
            'code': 200,
            'message': '添加成功',
            # 'result': data
        })

    @action(detail=True, methods=['get'])
    def test_file_analyze(self, request, pk=None):
        graph = self.get_object()
        # files = UnstructuredFile.objects.filter(graph=graph)
        #
        # # 组织生成文件信息
        # file_list = list()
        # for file_item in files:
        #     file_path = file_item.file_path
        #     file_name = file_item.file_name
        #     file_type = file_item.file_type
        #
        #     print("上传文件中...")
        #     # 文件上传
        #     file_id = dify_upload_file('非结构化处理', file_path, file_type)
        #     if file_id:
        #         file_list.append({
        #             'transfer_method': "local_file",
        #             'upload_file_id': f"{file_id}",
        #             'type': "document"
        #         })
        #     else:
        #         logger.error(f"文件上传失败, file_path:{file_path}")
        #
        #     break
        #
        # # 执行文件解析操作
        # # file_id = '8730c658-0fae-4368-b2dd-bbe0d7597e39'
        # print(file_list)
        # # 文件解析
        # result = dify_run_workflow('非结构化处理', {
        #     "file": file_list
        # })
        #
        # json_str = parse_workflow_response(result)
        # parse_result = json.loads(json_str)
        # print(parse_result)
        # total_rows = len(parse_result)
        # graph.build_total = total_rows
        # graph.save()
        process_unstructured_file.delay(26)
        return Response({
            'code': 200,
            'message': '上传成功',
            # 'file_id': file_id,
            # 'data': parse_result
        }, status=status.HTTP_200_OK)

    @action(detail=True, methods=['delete'])
    def delete(self, request, pk=None):
        # 删除构建图谱
        try:
            graph = self.get_object()  # 使用DRF提供的get_object方法来获取图谱实例
            graph.status = KnowledgeGraphStatus.DELETED.value
            graph.save()
            return Response({
                'code': 200,
                'message': '删除成功',
            }, status=status.HTTP_200_OK)
        except KnowledgeGraph.DoesNotExist:
            return Response(
                {'error': "未找到指定的图谱"},
                status=status.HTTP_404_NOT_FOUND
            )